import cv2
import os
from ultralytics import YOLO

def detect_person_in_video(video_path, output_dir, skip=5):
    # 创建输出目录（如果不存在的话）
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 提取文件路径中的目录和文件名
    directory, filename = os.path.split(video_path)  # 分离出目录和文件名
    name_without_extension = os.path.splitext(filename)[0]  # 去掉扩展名（.mp4）

    # 从路径中提取上一级目录的名称
    parent_dir = os.path.basename(os.path.dirname(video_path))

    # 组合新的文件名：G:/vlog/outputs/parent_dir/name_without_extension.jpg
    output_filename = os.path.join(output_dir, f"{parent_dir}/{name_without_extension}.jpg")

    # 创建子目录（如果不存在的话）
    output_subdir = os.path.dirname(output_filename)
    if not os.path.exists(output_subdir):
        os.makedirs(output_subdir)

    cap = cv2.VideoCapture(video_path)
    model = YOLO('yolov5su.pt')  # 加载YOLO模型

    frame_count = 0  # 帧计数器

    while True:
        ret, frame = cap.read()

        if not ret:
            break

        frame_count += 1  # 增加帧计数器

        # 每隔 skip 帧进行一次推理
        if frame_count % skip == 0:
            frame_resized = cv2.resize(frame, (640, 360))
            res = model(frame_resized, verbose=False)  # 设置verbose=False避免打印日志

            cls = res[0].boxes.cls  # 检测到的物体类型
            xyxy = res[0].boxes.xyxy  # 检测到的物体位置

            detect = False  # 是否检测到人

            # 遍历所有的检测结果，判断是否有检测到人
            for i, t in enumerate(cls):
                if t == 0:  # 0 表示 person 类型
                    detect = True
                    break

            if detect:
                # 如果检测到人，保存图像
                cv2.imwrite(output_filename, frame)  # 保存当前帧图像
                cap.release()
                return True  # 如果检测到人，返回 True

    cap.release()
    return False  # 如果视频中没有检测到人，返回 False


# 使用示例
video_path = "G:/vlog/182/2024-12-07-13-40-19.mp4"
output_dir = "G:/vlog_outputs"  # 指定保存图像的目录
skip = 5  # 每5帧进行一次YOLO推理

contains_person = detect_person_in_video(video_path, output_dir, skip)

if contains_person:
    print("视频中检测到人！")
else:
    print("视频中未检测到人。")
